'''
Password vault. Part of the cantankerous-contingency project

REQUIREMENTS: PyCrypto

INFO:Use a master password to manage your various passwords.

Each password managed by the vault has an:
 - Id
 - Value (the password itself)
 - Additional text information (could be the username associated with the password)

Although this module could be loaded into another program, it is intended as a
stand alone application. The reason being that once the master password has 
been entered, a cipher is created and stored in the PasswordVault object. 
Because of the way python works, external code can use it to read all of the 
passwords in the file.
'''

import os
import sys
import pickle
import base64
import getpass
import Crypto.Cipher.AES as AES
from urllib import urlopen

import helpgenerator

##########################################################
#### Main configuration
##########################################################

#Should be 16 (128 bit), 24 (192 bit) or 32 (256 bit)
#TODO: This should be user configurable
BLOCK_SIZE = 24

#Only used when called as a stand alone application.
DEFAULT_VAULT_LOCATION = os.environ["APPDATA"] + "\\pvdata.aes"

#---Other stuff. 
#Used for checking successful file decryption
FILE_HEADER = "[Cantankerous Contingency Password Vault]"

#---Try to use the GMZ module for backing up the passowrd vault file
USE_GMZ=True
if USE_GMZ:
    try:
        import GMZ
    except ImportError:
        #GMZ is not on a known path. Try a relative path to this file.
        try:
            sys.path.append(os.path.abspath(os.path.dirname(__file__) + "\\.."))
            import GMZ
        except ImportError:
            USE_GMZ=False
            print "NOT going to use GMZ to backup file"


##########################################################
#### Helper functions
##########################################################

#Simple search helper - pretty inefficent really
exists = lambda L, passid: True in map(lambda pw: pw.PasswordId == passid, L)

# These helper functions from http://www.codekoala.com/blog/2009/aes-encryption-python-using-pycrypto/
# 32 used since messages must be in blocks of 16 (ie, problems with block size of 24)
pad = lambda s: s + (32 - len(s) % 32) * " "
EncodeAES = lambda c, s: base64.b64encode(c.encrypt(pad(s)))
DecodeAES = lambda c, e: c.decrypt(base64.b64decode(e)).rstrip()

##########################################################
#### Classes
##########################################################

class PasswordInfo:
    def __init__(self, passid, value, info):
        self.PasswordId = passid
        self.PasswordValue = value
        self.Info = info

class PasswordVault:
    '''For interacting with decrypted passwords'''
    
    def __init__(self, accessCode, filePath):
        '''Create a password vault object. A new file is created if the given
        file path does not exist.
        
        accessCode: AKA Master Password
        filePath: Path to the file used to store the passwords.'''
        accessCode = str(accessCode)
        filePath = str(filePath)
        
        #TODO: Throw proper exception
        assert len(accessCode) == BLOCK_SIZE, "Wrong password length"
        
        self.haschanges = False #For backup of encrypted file
        self.changelist = "" #List of changes which have been made
        self.__filePath = filePath
        
        #Dont save the access code - directly create the cipher
        self.__cipher = AES.new(accessCode)
        
        self.__passwords = [] #list of PasswordInfo objects
        
        #Load passwords or create new password storage file.
        if os.path.isfile(filePath):
            self.__load();
        else:
            print "\nNOTE: Password file does not exist, going to create"
            self.__save();
    
    ######################################################
    #### High level access - callable from the interactive session
    ######################################################
    
    def list(self):
        '''Print a list of all password ids in the vault (but not the passwords
        themselves)
        help[1](List password ids)'''
        if len(self.__passwords) == 0:
            print "No passwords"
            return
        
        print ""
        for pw in self.__passwords:
            toprint = pw.PasswordId 
            if pw.Info != "":
                toprint += "   (" + pw.Info + ")"
            
            print toprint    
    
    def ls(self):
        '''help[2](List password ids \(same as list\))'''
        return self.list()


    def get(self, passId):
        '''help[3](Get the password for a given id)'''
        for pw in self.__passwords:
            if pw.PasswordId == passId:
                print ""
                print "Password: " + pw.PasswordValue
                if pw.Info != "":
                    print "Info: " + pw.Info
                return
        
        print "No password matching that id found"

    def new(self, passId, newValue=""):
        '''help[4](Create a new password entry with the given id)'''
        passId = str(passId)
        newValue = str(newValue)
        
        assert len(passId) > 0, "Must enter a password id"
        assert not exists(self.__passwords, passId), "Password id already in use"
        
        if newValue == "":
            newValue = getpass.getpass("New password:")
            if getpass.getpass("Enter again:") != newValue:
                print "Passwords do not match"
                return
        assert len(newValue) > 0, "Must enter a non-empty password"
        
        info = raw_input("Additional info (optional): ")
        
        self.__passwords.append(PasswordInfo(passId, newValue, info))
        
        #Sort the list
        self.__passwords.sort(key=lambda item: item.PasswordId)
        
        self.changelist += "\tNew password id: %s\r\n" % passId
        self.__save()


    def newgen(self, passId):
        '''help[5](Create a new password entry with the given id using a random passowrd)'''
        
        try:
            con = urlopen("https://www.dustbunnytycoon.com/pwgen")
            password = con.read()
            con.close()
        except IOError:
            print "Could not obtain password"
            return
        
        if len(password) == 0 or len(password) > 20:
            print "Problem with password"
            return
        
        print "Using randomly generated password: %s" % password
        
        self.new(passId, password)

    def delete(self, passId):
        '''help[6](Delete the given password id)'''
        passwd = None
        for pw in self.__passwords:
            if pw.PasswordId == passId:
                passwd = pw
                break
        
        if passwd == None:
            print "ID not found"
            return
        
        if raw_input("Are you sure you want to delete? [y/N]: ") == "y":
            self.__passwords.remove(passwd)
            
            self.changelist += "\tDeleted password id: %s\r\n" % passId
            self.__save()

    def change(self, passId):
        '''help[7](Change the password)'''
        passwd = None
        for pw in self.__passwords:
            if pw.PasswordId == passId:
                passwd = pw
                break
        
        if passwd == None:
            print "Password ID not found"
            return
        
        newValue = getpass.getpass("New password:")
        if getpass.getpass("Enter again:") != newValue:
            print "Passwords do not match"
            return
        
        assert len(newValue) > 0, "Must enter a non-empty password"
        pw.PasswordValue = newValue
        
        self.changelist += "\tChanged password: %s\r\n" % passId
        self.__save()


    def changeid(self, passId, newId):
        '''help[8](Change a password id)'''
        passId = str(passId)
        newId = str(newId)
        
        assert len(passId) > 0, "Must enter a password id"
        assert len(newId) > 0, "Must enter a new password id"
        assert not exists(self.__passwords, newId), "New password id already in use"
        
        passwd = None
        for pw in self.__passwords:
            if pw.PasswordId == passId:
                passwd = pw
                break
        
        pw.PasswordId = newId
        if passwd == None:
            print "Password ID not found"
            return
        
        #Sort the list
        self.__passwords.sort(key=lambda item: item.PasswordId)
        
        self.changelist += "\tChanged password id from \"" + passId + "\" to \"" + newId + "\"\r\n"
        self.__save()


    def changeinfo(self, passId):
        '''help[9](Update the additional information for the given password id)'''
        passwd = None
        for pw in self.__passwords:
            if pw.PasswordId == passId:
                passwd = pw
                break
        
        if passwd == None:
            print "ID not found"
            return
        
        print "Current info: " + passwd.Info
        passwd.Info = raw_input("New Info: ")
        
        self.changelist += "\tUpdated info for password id: %s\r\n" % passId
        self.__save()

    #### Internals
    
    def __load(self):
        '''Loads the current'''
        encryptedContentsFile = None
        try:
            encryptedContentsFile = open(self.__filePath, "rb")
            encryptedContents = encryptedContentsFile.read()
        finally:
            if not encryptedContentsFile is None:
                encryptedContentsFile.close()
        
        decryptedContents = DecodeAES(self.__cipher, encryptedContents)
        
        if not decryptedContents.startswith(FILE_HEADER):
            #TODO: Throw a proper exception
            assert False, "Decryption failed"
        
        decryptedContents = decryptedContents.replace(FILE_HEADER, "", 1)
        
        self.__passwords = pickle.loads(decryptedContents)
        assert type(self.__passwords) == type([]), "Depickle failed (?)"
    
    def __save(self):
        
        contents = pickle.dumps(self.__passwords)
        contents = FILE_HEADER + contents
        
        contents = EncodeAES(self.__cipher, contents)
        
        #Delete the current file - might want to change this since an error 
        #before save would cause loss of all data.
        if os.path.exists(self.__filePath):
            os.remove(self.__filePath)
        
        outputFile = None
        try:
            outputFile = open(self.__filePath, "wb")
            outputFile.write(contents)
        finally:
            outputFile.close()
        
        #If GMZ or some other backup mechanism is in place - indicate the file
        #should be backed up
        self.haschanges = True

##########################################################
#### Application
##########################################################

if __name__ == "__main__":
    print ""
    print "Starting password vault"
    try:
        pv = PasswordVault(getpass.getpass("Please enter access code (master password): "), DEFAULT_VAULT_LOCATION)
    except KeyboardInterrupt:
        print "\n\nAwww, please dont go!"
        exit()
    
    print "Decryption success"
    
    action = ""
    
    #Start the interactive session:
    while action.lower() != "exit":
        if action.startswith("_"):
            #Dont allow access to privates
            print "BAD User - Dont try that"
        elif action == "":
            pass
        elif action.lower() == "help":
            print helpgenerator.gethelp(pv)
        else:
            #A straight up split is not the best thing to do
            #It does not handle quote delimited arguments
            actionComponents = action.split()
            
            if len(actionComponents) == 0:
                pass
            
            methodname = actionComponents[0]
            if hasattr(pv, methodname): 
                _member = getattr(pv, methodname)
                try:
                    _member(*(actionComponents[1:]))
                except KeyboardInterrupt:
                    print "\nCanceled"
                except Exception as err:
                    print "ERROR: " + str(err)
                    #Dont exit
            else:
                print "Unknown command"
        
        print ""
        
        #Get the next command to run from the user. Commands are methods on the
        #PasswordVault object.
        try:
            action = raw_input(">>").strip()
        except KeyboardInterrupt:
            action = ""
            print "Please type \'exit\' to quit"
    
    if USE_GMZ and pv.haschanges:
        print ""
        print "Backing up password vault file using GMZ"
        
        #TODO: I suppose network connectivity should be checked before calling
        GMZ.zipandsend( DEFAULT_VAULT_LOCATION, 
                        "Password Vault Data (Encrypted)", 
                        pv.changelist)
    
    print "Quitting"
    
else:
    print __name__
    print "WARNING: Loading this file as a module is NOT recommended for security reasons"